# -*- coding: utf-8 -*-
# ☯ Author: ChinaPython
# ☯ Email : chinapython@yeah.net
# ☯ Date  : 2021/11/3 14:37
import re

import cv2
import requests
import numpy as np
import paddlehub as hub


class CharRec:
    def __init__(self):
        self.ocr = hub.Module(name="chinese_ocr_db_crnn_mobile")

    def distinguish(self, np_image):
        return self.ocr.recognize_text(
            images=[np_image],  # 图片数据，ndarray.shape 为 [H, W, C]，BGR格式；
            use_gpu=False,  # 是否使用 GPU；若使用GPU，请先设置CUDA_VISIBLE_DEVICES环境变量
            output_dir=r"C:\test",  # 图片的保存路径，默认设为 ocr_result；
            visualization=True,  # 是否将识别结果保存为图片文件；
            box_thresh=0.5,  # 检测文本框置信度的阈值；
            text_thresh=0.5  # 识别中文文本置信度的阈值；
        )

    @staticmethod
    def download(url):
        try:
            image = requests.get(url).content
            if b"JFIF" not in image and b"PNG" not in image and b"BM" not in image:
                return
            return cv2.imdecode(np.fromstring(image, np.uint8), 1)
        except Exception as e:
            return print(e)

    def image_detection(self, url):
        content = self.download(url)
        if content is None:
            return print("image download failed")
        try:
            items = self.distinguish(content)
            return "/".join(re.findall("'text': '(.*?)',", str(items)))
        except Exception as e:
            return print(e)


if __name__ == '__main__':
    start = CharRec()
    start.image_detection("https://www.aicloud.com:8889/web/img/banner/15535858535_PC.jpg")
